#-*- coding: utf-8 -*-

#
# simplified python interface to JSBsim
#

import math
import numpy as np
import matplotlib.pyplot as plt

import pat.dynamic_model as dm
import pat.frames as fr
import pat.utils as pu

class DynamicModel(dm.DynamicModel):

    sv_x    = 0
    sv_y    = 1
    sv_z    = 2
    sv_xd   = 3
    sv_yd   = 4
    sv_zd   = 5
    sv_phi  = 6
    sv_theta = 7
    sv_psi  = 8
    sv_p    = 10
    sv_q    = 11
    sv_r    = 12
    sv_size = 13

    iv_da   = 0
    iv_de   = 1
    iv_dr   = 2
    iv_th   = 3
    iv_size = 4

    
    def __init__(self, param):
        dm.DynamicModel.__init__(self)
        self.X = np.zeros(DynamicModel.sv_size)
        try:
#            self.jsbsim = __import__("JSBSim")
            self.jsbsim = __import__("libpat_dm_jsbsim")
            self.fdex = self.jsbsim.FGFDMExec( )
            self.fdex.SetDebugLevel( 0 )
            self.fdex.SetAircraftPath('/usr/share/games/flightgear/Aircraft')
            self.fdex.LoadModel('c172p')
            propulsion = self.fdex.GetPropulsion( )
            fcs = self.fdex.GetFCS( )
            fcs.SetThrottleCmd( 0, 0.25 )
            fcs.SetMixtureCmd( 0, 0.87 )
            propulsion.SetMagnetos( 3 )
            propulsion.SetStarter( 1 )
            self.reset()
        except ImportError:
            print "Error: could not load native JSBsim  library"
            

    def name(self):
        return "fixed wing jsbsim" 

    def reset(self, X0=None):
        if X0<>None: self.X = X0
        else: self.X = np.zeros(DynamicModel.sv_size)
        fdex = self.fdex
        ic = self.fdex.GetIC( )
        ic.SetLatitudeDegIC( 47.0 ) ; ic.SetLongitudeDegIC( 122.0 ); ic.SetAltitudeFtIC( 1500. )
        ic.SetUBodyFpsIC( 70.0 ); ic.SetVBodyFpsIC( 0.0 ) ;ic.SetWBodyFpsIC( 0.0 )
        ic.SetPhiDegIC( 0.0 ); ic.SetThetaDegIC( 0.0 ); ic.SetPsiDegIC( 0.0 )
        self.fdex.RunIC( )
        self.fdex.DoTrim( self.jsbsim.tGround )
        self.read_jsbsim()


    def trim(self, args=None):
        Xe = np.zeros(DynamicModel.sv_size)
        Ue = np.array([-0.00033625,  0.00438276,  0.16024332,  0.10031169])
        return Xe, Ue

    def read_jsbsim(self):
        propagate = self.fdex.GetPropagate()
        location = propagate.GetLocation( )
        self.X[self.sv_x], self.X[self.sv_y], self.X[self.sv_z] = location.GetEntries( )
        self.X[self.sv_xd], self.X[self.sv_yd], self.X[self.sv_zd] = propagate.GetUVW()
        self.X[self.sv_p], self.X[self.sv_q], self.X[self.sv_r] = propagate.GetPQR()
        self.X[self.sv_phi], self.X[self.sv_theta], self.X[self.sv_psi] = propagate.GetEuler( )
        

    def run(self, dt, U):
#        fcs = self.jsbsim.FGFDMExec().GetFCS()
        fcs = self.fdex.GetFCS()
        fcs.SetDaCmd( U[self.iv_da] )  # aileron
        fcs.SetDeCmd( U[self.iv_de] )  # elevator
        fcs.SetDrCmd( U[self.iv_dr] )  # rudder
        fcs.SetThrottleCmd( 0, U[0] )  # throttle
        self.fdex.Run()
        self.read_jsbsim()
        return self.X

    def param(self):
        return str(self.p)

    def state_SixDOFfEuclidianEuler(self):
        X = np.zeros(fr.SixDOFfEuclidianEuler.size)
        X[fr.SixDOFfEuclidianEuler.x:fr.SixDOFfEuclidianEuler.z+1] = self.X[self.sv_x:self.sv_z+1]
#        X[fr.SixDOFfEuclidianEuler.phi:fr.SixDOFfEuclidianEuler.r+1] = self.X[self.sv_phi:self.sv_r+1
        X[fr.SixDOFfEuclidianEuler.phi:fr.SixDOFfEuclidianEuler.psi+1] = self.X[self.sv_phi:self.sv_psi+1]
        return X

    def state_str(self):
        return """pos: {:-.2f}, {:-.2f}, {:-.2f} m
        vel: {:-.2f}, {:-.2f}, {:-.2f} m/s
        att:    {:-.2f}, {:-.2f}, {:-.2f} deg
        """.format(self.X[self.sv_x], self.X[self.sv_y], self.X[self.sv_z],
                   self.X[self.sv_xd], self.X[self.sv_yd], self.X[self.sv_zd],
                   pu.deg_of_rad(self.X[self.sv_phi]), pu.deg_of_rad(self.X[self.sv_theta]), pu.deg_of_rad(self.X[self.sv_psi]))


    def plot_state(self, time, X):

        ax = plt.subplot(4,3,1)         # x
        plt.plot(time, X[:,self.sv_x])
        pu.decorate(ax, title="$x$")
        ax=plt.subplot(4,3,2)           # y
        plt.plot(time, X[:,self.sv_y])
        pu.decorate(ax, title="$y$")
        ax=plt.subplot(4,3,3)           # z
        plt.plot(time, X[:,self.sv_z]) 
        pu.decorate(ax, title="$z$")

        ax = plt.subplot(4,3,4)         # xd
        plt.plot(time, X[:,self.sv_xd])
        pu.decorate(ax, title="$\dot{x}$")
        ax=plt.subplot(4,3,5)           # yd
        plt.plot(time, X[:,self.sv_yd])
        pu.decorate(ax, title="$\dot{y}$")
        ax=plt.subplot(4,3,6)           # zd
        plt.plot(time, X[:,self.sv_yd]) 
        pu.decorate(ax, title="$\dot{z}$")

        ax=plt.subplot(4,3,7)
        plt.plot(time, pu.deg_of_rad(X[:,self.sv_phi]))  # phi
        pu.decorate(ax, title="$\phi$")
        ax=plt.subplot(4,3,8)
        plt.plot(time, pu.deg_of_rad(X[:,self.sv_theta])) # theta
        pu.decorate(ax, title="$\\theta$")
        ax=plt.subplot(4,3,9)
        plt.plot(time, pu.deg_of_rad(X[:,self.sv_psi])) # psi
        pu.decorate(ax, title="$\psi$")
        
        ax=plt.subplot(4,3,10)
        plt.plot(time, pu.deg_of_rad(X[:,self.sv_p]))  # p
        pu.decorate(ax, title="$p$")
        ax=plt.subplot(4,3,11)
        plt.plot(time, pu.deg_of_rad(X[:,self.sv_q])) # q
        pu.decorate(ax, title="$q$")
        ax=plt.subplot(4,3,12)
        plt.plot(time, pu.deg_of_rad(X[:,self.sv_r])) # r
        pu.decorate(ax, title="$r$")

    def plot_state_alt(self, time, X):
        pass
    
